import torch
import sys
import importlib.util

# 检查Python版本
def check_python_version():
    """检查Python版本是否满足要求"""
    required_version = (3, 8)
    current_version = sys.version_info
    print(f"Python版本: {current_version.major}.{current_version.minor}.{current_version.micro}")
    if current_version < required_version:
        print(f"错误: Python版本至少需要{required_version[0]}.{required_version[1]}，当前版本为{current_version.major}.{current_version.minor}")
        return False
    else:
        print("Python版本检查通过")
        return True

# 检查必要的库
def check_required_libraries():
    """检查必要的库是否已安装"""
    required_libraries = {
        'torch': '>=2.6',
        'transformers': '>=4.53.0',
        'PIL': None,
        'cv2': None,
        'json': None
    }
    
    missing_libraries = []
    for lib, version in required_libraries.items():
        try:
            if lib == 'PIL':
                importlib.import_module('PIL')
                import PIL
                print(f"Pillow版本: {PIL.__version__}")
            elif lib == 'cv2':
                importlib.import_module('cv2')
                import cv2
                print(f"OpenCV版本: {cv2.__version__}")
            elif lib == 'torch':
                importlib.import_module('torch')
                import torch
                print(f"PyTorch版本: {torch.__version__}")
                if version and torch.__version__ < version[2:]:  # 去掉'>='
                    print(f"警告: PyTorch版本建议至少为{version[2:]}，当前版本为{torch.__version__}")
            elif lib == 'transformers':
                importlib.import_module('transformers')
                import transformers
                print(f"Transformers版本: {transformers.__version__}")
                if version and transformers.__version__ < version[2:]:  # 去掉'>='
                    print(f"警告: Transformers版本建议至少为{version[2:]}，当前版本为{transformers.__version__}")
            else:
                importlib.import_module(lib)
                print(f"{lib}库已安装")
        except ImportError:
            missing_libraries.append(f"{lib}{version if version else ''}")
    
    if missing_libraries:
        print(f"错误: 缺少以下必要库: {', '.join(missing_libraries)}")
        print("请运行: pip install -r requirements.txt 来安装缺少的依赖")
        return False
    else:
        print("所有必要的库检查通过")
        return True

# 检查GPU环境
def check_gpu_environment():
    """检查GPU环境是否可用"""
    try:
        import torch
        print(f"CUDA可用: {torch.cuda.is_available()}")
        if torch.cuda.is_available():
            print(f"CUDA版本: {torch.version.cuda}")
            print(f"GPU名称: {torch.cuda.get_device_name(0)}")
            print(f"GPU内存: {torch.cuda.get_device_properties(0).total_memory / 1024**3:.2f} GB")
            return True
        else:
            print("警告: 未检测到可用的GPU，程序将使用CPU运行，这会导致推理速度较慢")
            print("请确保已正确安装CUDA和cuDNN，并且显卡支持CUDA")
            return False
    except ImportError:
        print("错误: 无法检查GPU环境，因为PyTorch未安装")
        return False

# 主函数
def main():
    """主函数，运行所有检查"""
    print("\n===== MiniCPM-V 环境检查 =====")
    
    python_ok = check_python_version()
    print()
    libraries_ok = check_required_libraries()
    print()
    gpu_ok = check_gpu_environment()
    
    print("\n===== 环境检查总结 =====")
    
    if python_ok and libraries_ok:
        if gpu_ok:
            print("环境检查通过！您的系统满足运行MiniCPM-V的所有要求。")
            print("建议: 您可以使用GPU加速，推理性能会更好。")
        else:
            print("环境检查基本通过，但建议配置GPU环境以获得更好的性能。")
        return 0
    else:
        print("环境检查未通过，请根据上述提示解决问题后再运行程序。")
        return 1

if __name__ == "__main__":
    sys.exit(main())